import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import TransformerEncoder, TransformerEncoderLayer
import math

class TextHybridModel(nn.Module):
    """融合CNN与Transformer的文本分类模型"""
    def __init__(self, 
                 vocab_size,        # 词汇表大小
                 embed_dim=128,     # 词嵌入维度
                 hidden_dim=256,    # 隐藏层维度
                 num_classes=2,     # 分类类别数
                 num_heads=4,       # Transformer注意力头数
                 num_layers=2,      # Transformer编码器层数
                 cnn_kernel_sizes=[3,5],  # CNN卷积核尺寸
                 cnn_out_channels=64,     # CNN输出通道数
                 max_seq_len=128,    # 最大序列长度
                 dropout=0.1):       # dropout概率
        super().__init__()
        
        # 1. 词嵌入层
        self.embedding = nn.Embedding(
            num_embeddings=vocab_size,
            embedding_dim=embed_dim,
            padding_idx=0  # 假设0为填充符
        )
        
        # 2. 位置编码（用于Transformer）
        self.pos_encoder = PositionalEncoding(
            d_model=embed_dim,
            dropout=dropout,
            max_len=max_seq_len
        )
        
        # 3. CNN层（多尺度卷积提取局部特征）
        self.cnn_layers = nn.ModuleList([
            nn.Conv1d(
                in_channels=embed_dim,
                out_channels=cnn_out_channels,
                kernel_size=k,
                padding=k//2  # 保持序列长度不变
            ) for k in cnn_kernel_sizes
        ])
        # CNN输出特征维度（多尺度拼接）
        self.cnn_total_dim = cnn_out_channels * len(cnn_kernel_sizes)
        
        # 4. 特征融合层（将CNN特征与嵌入特征合并）
        self.fusion_proj = nn.Linear(
            in_features=embed_dim + self.cnn_total_dim,
            out_features=hidden_dim
        )
        
        # 5. Transformer编码器
        transformer_layer = TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=num_heads,
            dim_feedforward=hidden_dim*4,  # 前馈网络维度
            dropout=dropout,
            batch_first=True  # 批次维度在前
        )
        self.transformer_encoder = TransformerEncoder(
            encoder_layer=transformer_layer,
            num_layers=num_layers
        )
        
        # 6. 分类头
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim//2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim//2, num_classes)
        )
        
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        """
        Args:
            x: 输入序列，形状 [batch_size, seq_len]
            mask: 掩码（可选），形状 [batch_size, seq_len]，1表示有效位置
        Returns:
            logits: 分类输出，形状 [batch_size, num_classes]
        """
        # 1. 词嵌入 + 位置编码
        embed = self.embedding(x)  # [batch, seq_len, embed_dim]
        embed = self.pos_encoder(embed)  # 加入位置信息
        
        # 2. CNN特征提取（需调整维度为[batch, embed_dim, seq_len]）
        cnn_input = embed.permute(0, 2, 1)  # 调整通道维度位置
        cnn_features = []
        for conv in self.cnn_layers:
            feat = conv(cnn_input)  # [batch, cnn_out_channels, seq_len]
            feat = F.relu(feat)
            cnn_features.append(feat.permute(0, 2, 1))  # 恢复为[batch, seq_len, cnn_out]
        cnn_cat = torch.cat(cnn_features, dim=-1)  # [batch, seq_len, cnn_total_dim]
        
        # 3. 特征融合（嵌入特征 + CNN特征）
        fused = torch.cat([embed, cnn_cat], dim=-1)  # [batch, seq_len, embed_dim + cnn_total]
        fused = self.fusion_proj(fused)  # [batch, seq_len, hidden_dim]
        fused = self.dropout(fused)
        
        # 4. Transformer编码
        if mask is not None:
            # 转换为Transformer需要的掩码格式（True表示忽略）
            attn_mask = (mask == 0).unsqueeze(1).repeat(1, fused.shape[1], 1)  # [batch, seq_len, seq_len]
        else:
            attn_mask = None
        transformer_out = self.transformer_encoder(fused, src_mask=attn_mask)  # [batch, seq_len, hidden_dim]
        
        # 5. 聚合序列特征（取CLS标记或全局平均）
        # 这里使用全局平均池化
        seq_feat = transformer_out.mean(dim=1)  # [batch, hidden_dim]
        
        # 6. 分类输出
        logits = self.classifier(seq_feat)  # [batch, num_classes]
        return logits


class PositionalEncoding(nn.Module):
    """位置编码模块（Transformer必备）"""
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        # 预计算位置编码
        position = torch.arange(max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
        pe = torch.zeros(max_len, 1, d_model)
        pe[:, 0, 0::2] = torch.sin(position * div_term)  # 偶数维度用正弦
        pe[:, 0, 1::2] = torch.cos(position * div_term)  # 奇数维度用余弦
        self.register_buffer('pe', pe)  # 不参与训练的参数

    def forward(self, x):
        """
        Args:
            x: 嵌入特征，形状 [seq_len, batch_size, d_model] 或 [batch_size, seq_len, d_model]
        Returns:
            加入位置编码的特征
        """
        # 适配batch_first格式
        if x.ndim == 3 and x.shape[1] == x.size(1):  # [batch, seq, dim]
            x = x + self.pe[:x.size(1)].permute(1, 0, 2)  # 调整pe维度为[1, seq_len, dim]
        else:
            x = x + self.pe[:x.size(0)]
        return self.dropout(x)


# 测试代码
if __name__ == "__main__":
    # 超参数设置
    vocab_size = 10000
    batch_size = 8
    seq_len = 32
    
    # 初始化模型
    model = TextHybridModel(
        vocab_size=vocab_size,
        embed_dim=128,
        hidden_dim=256,
        num_classes=3,
        num_heads=4,
        num_layers=2
    )
    
    # 生成随机输入
    x = torch.randint(0, vocab_size, (batch_size, seq_len))  # 随机序列
    mask = torch.ones(batch_size, seq_len)  # 掩码（全为有效）
    
    # 前向传播
    logits = model(x, mask)
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {logits.shape}")  # 应输出 [8, 3]